"""
数据获取模块
负责从不同数据源获取股票数据
"""

import yfinance as yf
import pandas as pd
import numpy as np
from typing import Dict, List, Optional, Tuple
from datetime import datetime, timedelta
import time


class DataFetcher:
    """数据获取器"""
    
    def __init__(self):
        """初始化数据获取器"""
        self.cache = {}  # 数据缓存
        self.last_request_time = 0
        self.min_request_interval = 0.1  # 最小请求间隔（秒）
    
    def fetch_stock_data(self, symbol: str, period: str = "1y", 
                        interval: str = "1d") -> Optional[pd.DataFrame]:
        """
        获取股票数据
        
        Args:
            symbol: 股票代码
            period: 时间周期 ('1d', '5d', '1mo', '3mo', '6mo', '1y', '2y', '5y', '10y', 'ytd', 'max')
            interval: 数据间隔 ('1m', '2m', '5m', '15m', '30m', '60m', '90m', '1h', '1d', '5d', '1wk', '1mo', '3mo')
            
        Returns:
            Optional[pd.DataFrame]: 股票数据
        """
        # 检查缓存
        cache_key = f"{symbol}_{period}_{interval}"
        if cache_key in self.cache:
            cached_data, timestamp = self.cache[cache_key]
            # 缓存有效期1小时
            if time.time() - timestamp < 3600:
                return cached_data
        
        # 控制请求频率
        self._rate_limit()
        
        try:
            # 获取数据
            ticker = yf.Ticker(symbol)
            data = ticker.history(period=period, interval=interval)
            
            if data.empty:
                return None
            
            # 清理数据
            data = data.dropna()
            
            # 缓存数据
            self.cache[cache_key] = (data, time.time())
            
            return data
            
        except Exception as e:
            print(f"获取 {symbol} 数据时出错: {e}")
            return None
    
    def fetch_multiple_stocks(self, symbols: List[str], period: str = "1y", 
                            interval: str = "1d") -> Dict[str, pd.DataFrame]:
        """
        获取多个股票数据
        
        Args:
            symbols: 股票代码列表
            period: 时间周期
            interval: 数据间隔
            
        Returns:
            Dict: 股票数据字典
        """
        results = {}
        
        for symbol in symbols:
            data = self.fetch_stock_data(symbol, period, interval)
            if data is not None:
                results[symbol] = data
            
            # 在请求之间添加延迟
            time.sleep(self.min_request_interval)
            
        return results
    
    def _rate_limit(self):
        """速率限制"""
        current_time = time.time()
        time_since_last_request = current_time - self.last_request_time
        
        if time_since_last_request < self.min_request_interval:
            sleep_time = self.min_request_interval - time_since_last_request
            time.sleep(sleep_time)
            
        self.last_request_time = time.time()
    
    def get_stock_info(self, symbol: str) -> Dict:
        """
        获取股票基本信息
        
        Args:
            symbol: 股票代码
            
        Returns:
            Dict: 股票信息
        """
        self._rate_limit()
        
        try:
            ticker = yf.Ticker(symbol)
            info = ticker.info
            return {
                'symbol': symbol,
                'name': info.get('longName', 'N/A'),
                'sector': info.get('sector', 'N/A'),
                'industry': info.get('industry', 'N/A'),
                'market_cap': info.get('marketCap', 0),
                'pe_ratio': info.get('trailingPE', 0),
                'dividend_yield': info.get('dividendYield', 0)
            }
        except Exception as e:
            print(f"获取 {symbol} 信息时出错: {e}")
            return {
                'symbol': symbol,
                'name': 'N/A',
                'sector': 'N/A',
                'industry': 'N/A',
                'market_cap': 0,
                'pe_ratio': 0,
                'dividend_yield': 0
            }


def main():
    """主函数 - 数据获取演示"""
    print("数据获取模块演示")
    print("="*30)
    
    # 创建数据获取器
    fetcher = DataFetcher()
    
    # 获取单个股票数据
    print("\n获取AAPL股票数据...")
    data = fetcher.fetch_stock_data("AAPL", period="1mo")
    if data is not None:
        print(f"获取到 {len(data)} 条数据")
        print(data.tail())
    
    # 获取股票信息
    print("\n获取AAPL股票信息...")
    info = fetcher.get_stock_info("AAPL")
    for key, value in info.items():
        print(f"{key}: {value}")

if __name__ == "__main__":
    main()